import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import random
import time


class CustomBN2D(nn.Module):

    def __init__(self, num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True):
        super().__init__()
        self.eps = eps
        self.momentum = momentum
        self.affine = affine
        if self.affine:
            self.weight = nn.Parameter(torch.ones(1, num_features, 1, 1))
            self.bias = nn.Parameter(torch.zeros(1, num_features, 1, 1))
        else:
            self.weight = None
            self.bias = None
        self.track_running_stats = track_running_stats
        self.total_batch_mean = torch.ones(num_features)
        self.total_batch_var = torch.zeros(num_features)
        self.training = True

    def get_mean_var(self, x):

        if self.training:
            b, c, h, w = x.shape
            x_trans = x.permute(1, 0, 2, 3)
            curr_batch_mean = x_trans.contiguous().view(c, -1).mean(dim=-1)
            curr_batch_var = x_trans.contiguous().view(c, -1).var(dim=-1, unbiased=False)
            x_hat = (x_trans.T - curr_batch_mean).div(torch.sqrt(curr_batch_var + self.eps)).T
            x_hat = x_hat.permute(1, 0, 2, 3)
        else:
            curr_batch_mean, curr_batch_var = self.total_batch_mean, self.total_batch_var
            x_hat = (x - curr_batch_mean).div(torch.sqrt(curr_batch_var + self.eps))

        if self.affine:
            transout = self.weight * x_hat + self.bias
        else:
            transout = x_hat
        return curr_batch_mean, curr_batch_var, transout

    def running_update(self, curr_batch_mean, curr_batch_var):
        if self.training:
            self.total_batch_mean = (1 - self.momentum) * curr_batch_mean + self.momentum * self.total_batch_mean
            self.total_batch_var = (1 - self.momentum) * curr_batch_var + self.momentum * self.total_batch_var

    def forward(self, x):
        # x shape : b,c,f
        assert len(x.shape) == 4, ValueError("input shape is error , you should input the x in the shape of 4")
        curr_batch_mean, curr_batch_var, transout = self.get_mean_var(x)
        self.running_update(curr_batch_mean, curr_batch_var)
        return transout


def judgeBN(x_input):
    c=x_input.shape[1]
    bn2d = nn.BatchNorm2d(c)
    x_output1 = bn2d(x_input)
    my_bn = CustomBN2D(c)
    x_output2 = my_bn(x_input)
    my_bn1 = CustomBN2D(c,affine=False)
    x_output3 = my_bn1(x_input)
    print(x_output1.shape)
    print(x_output2.shape)
    print(x_output3.shape)


x = torch.rand(2,3,4,5)
judgeBN(x)